package dev.yorkecao.demo.mybatis.aop;

import dev.yorkecao.demo.mybatis.annotation.DynamicDataSource;
import dev.yorkecao.demo.mybatis.config.DemoConfig;
import dev.yorkecao.demo.mybatis.config.MultipleDataSource;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@Aspect
@Component
public class DynamicDataSourceAspect {

    @Autowired
    private DemoConfig demoConfig;

    @Pointcut("@annotation(dev.yorkecao.demo.mybatis.annotation.DynamicDataSource)")
    public void mapperMethod() {}

    @Around("mapperMethod()")
    public Object setDataSource(ProceedingJoinPoint joinPoint) throws Throwable {
        Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
        DynamicDataSource dynamicDataSource = method.getAnnotation(DynamicDataSource.class);
        String[] configArray = dynamicDataSource.value();
        Map<String, String> configMap = Arrays.stream(configArray)
                .map(t -> t.split(":"))
                .collect(Collectors.toMap(t -> t[0], t -> t[1]));
        String dataSourceKey = configMap.get(demoConfig.getEnv());
        log.debug("设置 dataSource：{}", dataSourceKey);
        MultipleDataSource.setCurrentLookupKey(dataSourceKey);

        try {
            return joinPoint.proceed();
        } finally {
            log.info("移除 dataSource：{}", dataSourceKey);
            MultipleDataSource.clearLookupKey();
        }
    }
}
